##########################################################################################

library(dplyr)
library(ggplot2)
library(data.table)
library(RColorBrewer)
library(optparse)

##########################################################################################

option_list <- list(
    make_option(c("--mut_rate_gene_file"), type = "character") ,
    make_option(c("--smg_file"), type = "character") ,
    make_option(c("--images_path"), type = "character")
)

if(1!=1){

    work_dir <- "~/20220915_gastric_multiple/dna_combinePublic/"
    mut_rate_gene_file <- "MutRate_baline_Tobacco.tsv"
    smg_file <- "All_driver.list"

}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

mut_rate_gene_file <- opt$mut_rate_gene_file
smg_file <- opt$smg_file
images_path <- opt$images_path


dir.create(images_path , recursive = T)

###########################################################################################

col <- c("#ecac54","#1e69b4")
names(col) <- c("Ever" , "Never")

###########################################################################################

dat_mutRateGene <- data.frame(fread(mut_rate_gene_file))
smg <- data.frame(fread(smg_file))$Gene_Symbol

###########################################################################################

dat_mutRateGene <- subset(dat_mutRateGene , Hugo_Symbol %in% smg)

dat_plot <- rbind( dat_mutRateGene )
dat_plot$Class <- factor( dat_plot$Class , levels = c("IM" , "GC" , "IGC" , "DGC") , order = T )
dat_plot$value_text <-round(dat_plot$MutRate , 2) * 100

###########################################################################################

dat_plot <- subset( dat_plot , From == "All" & Tobacco != "unknown")

for(type in c("MSS","CIN","GS")){
  dir.create(paste0(images_path,"/",type) , recursive = T)
  dat_plot_1 <- subset(dat_plot,Molecular.subtype==type)
  ##后续拆分出来IM,IGC,DGC
  for(class in c("IM","IGC","DGC")){
    print(class)
    dat_plot_use <- subset( dat_plot_1 , Class == class)
    ever_num <- unique(subset(dat_plot_use , Tobacco=="Smoke")$SampleNum)
    never_num <- unique(subset(dat_plot_use , Tobacco=="No")$SampleNum)
    dat_plot_use$Tobacco <- ifelse(dat_plot_use$Tobacco=="Smoke","Ever","Never")
    ## 计算P值
    dat_plot_use$p.value = ""
    dat_plot_use$p_text = ""
    
    result_tmp <- c()
    
    for(geneN in unique(dat_plot_use$Hugo_Symbol)){
      
      print(geneN)
      
      tmp_1 <- subset( dat_plot_use , Hugo_Symbol == geneN &  Tobacco %in% c("Ever") )
      tmp_2 <- subset( dat_plot_use , Hugo_Symbol == geneN &  Tobacco %in% c("Never") )
      
      if(nrow(tmp_1)==0){
        tmp_1 <- tmp_2
        tmp_1$Tobacco <- "Ever"
        tmp_1$SampleNum <- ever_num
        tmp_1$MutNum <- 0
        tmp_1$MutRate <- 0
        tmp_1$value_text <- "0"
      }
      
      if(nrow(tmp_2)==0){
        tmp_2 <- tmp_1
        tmp_2$Tobacco <- "Never"
        tmp_2$SampleNum <- never_num
        tmp_2$MutNum <- 0
        tmp_2$MutRate <- 0
        tmp_2$value_text <- "0"
      }
      
      tmp <- rbind( tmp_1 , tmp_2 )
      
      tmp_fisher <- matrix(c(tmp$MutNum , tmp$SampleNum - tmp$MutNum) , ncol = 2)
      p <- fisher.test(tmp_fisher)$p.value
      
      tmp$p.value <- p
      
      result_tmp <- rbind( result_tmp , tmp )
      
    }
    
    gene_order <- c("TP53","ARID1A","CDH1","APC","SMAD4","MUC6","PIK3CA",
                    "CTNNB1","RHOA","ERBB2","CFTR","KRAS","MAP2K7","ARID2",
                    "RNF43","TGFBR2","BMP6","FBXW7","CDKN2A","MTRR")
    result_tmp$Hugo_Symbol <- factor( result_tmp$Hugo_Symbol , 
                                      levels = gene_order , order = T)
    
    
    result_use <- result_tmp
    
    result_use$p <- result_use$p.value
    result_use$p_text=ifelse(result_use$p>=0.05,"","*")
    result_use$p_text=ifelse(result_use$p<0.05 & result_use$p>0.01,"*",result_use$p_text)
    result_use$p_text=ifelse(result_use$p<0.01 & result_use$p>0.001,"**",result_use$p_text)
    result_use$p_text=ifelse(result_use$p<0.001 ,"***",result_use$p_text)
    
    result_use$p_pos <- 0.8
    result_use$p_pos <- as.numeric(result_use$p_pos)
    
    
    result_use$percent_pos <- result_use$MutRate + 0.01
    result_use$percent_pos <- as.numeric(result_use$percent_pos)
    
    #print(result_use)
    
    p <- ggplot(result_use,mapping = aes(x = Hugo_Symbol , y = MutRate , fill = Tobacco)) +
      geom_bar(stat = 'identity', position = 'dodge' , width = 0.8 , color = 'black') + 
      #facet_grid(vars(Type) , scales = "free")+
      theme_bw() +
      scale_y_continuous(
        breaks = seq(-0.8,0.8,0.2) , 
        label = c( "80%" , "60%" , "40%" , "20%" , 
                   0 ,
                   "20%" , "40%" , "60%" , "80%"
        )
      ) +
      geom_text(aes(label= value_text , x = Hugo_Symbol, y = percent_pos ), position=position_dodge(0.9),size=5, vjust=0 ,color="black", face='bold' , family="Helvetica")+
      geom_text(aes(label=p_text , x = Hugo_Symbol , y = p_pos ),size=7,family="Helvetica")+
      #geom_text(aes(label=Type , x = 18 , y = label_pos ),size=5,family="Helvetica")+
      xlab("") +
      ylab('Percentage of samples with mutations') +
      scale_fill_manual(values = col) +
      theme(
        title =element_text(size=4, face='bold'),
        legend.title = element_blank(),
        legend.text = element_text(size = 10),
        legend.key.width = unit(1, "cm"),
        legend.key.height = unit(1, "cm"),
        legend.position = c(0.90,0.8) ,
        strip.text = element_blank(),
        axis.text.x = element_text(angle = 45, vjust = 1, hjust = 1 , size = 14  , color = 'black' , family="Helvetica" ) ,
        axis.title.y =  element_text(size = 16 , color = 'black' ) ,
        axis.text.y =  element_text(size = 14 , color = 'black' , family="Helvetica") ,
        axis.ticks.length = unit(0.2, "cm") ,
        panel.grid=element_blank() 
      )
    
    width <- 14
    height <- 5.3
    images_name <- paste0(images_path ,"/",type, "/MutRate.compare.",class,".Tobacco.Gene.pdf")
    ggsave( images_name , p , width = width , height = height )
    images_name <- paste0(images_path,"/",type , "/MutRate.compare.",class,".Tobacco.Gene.tsv")
    write.table( result_use , images_name , row.names = F , sep = "\t" , quote = F )
  }
  
}



